import os
import numpy as np
import numpy.fft as fft
import pylab as plt
import scipy.signal as sg


class GPR:
    
    def __init__(self):
        pass
    
    
    def loadrd3(self, filename):
        headerfile = filename+'.rad'
        datafile = filename+'.rd3'
        self.filename = os.path.basename(filename)
        # Read the header file
        fH = open(headerfile,'r')
        sH = fH.readlines()
        fH.close()
        
        self.samples = int(sH[0].split(':')[1].split('\\')[0])
        self.traces = int(sH[22].split(':')[1].split('\\')[0])
        self.timewindow = float(sH[18].split(':')[1].split('\\')[0]) # in ns
        self.deltaT = (self.timewindow*1e-9)/(self.samples-1)
        self.t = np.arange(self.samples)*self.deltaT
        print 'samples: ', self.samples
        print 'traces: ', self.traces
        print 'time window: ', self.timewindow
        print 'delta T: ', self.deltaT
        print self.t
            
        # Read the data file
        fD = open(datafile, 'rb')
        sD = fD.read()
        fD.close()
        data = np.fromstring(sD, 'int16')
        B = np.reshape(data, (self.traces, self.samples))
        Xaxis = np.arange(self.samples)
        Yaxis = np.arange(self.traces)
        self.Tdata = np.transpose(B)
        
        
    def subtractDC(self):
        Nfft = 2048
        for i in range(self.traces):
            f = fft.fft(self.Tdata[:,i], Nfft)
            f[0] = 0.0 + 0.0j
            self.Tdata[:,i] = fft.ifft(f, Nfft)
#        freqs = (np.arange(Nfft) / (self.deltaT * Nfft))
#        for i in range(Nfft/2):
#            freqs[Nfft/2+i] = - freqs[Nfft/2-i]
#        freqs = freqs[0:Nfft/2]
#        
#        plt.plot(trace, self.t, 'b', trace_f, self.t, 'r')
#        a = plt.axis()
#        plt.axis([a[0], a[1], a[3], a[2]])
#        plt.figure()
#        plt.plot(freqs, np.abs(np.real(f[0:Nfft/2])))
#        plt.figure()
#        plt.plot(freqs, np.abs(np.imag(f[0:Nfft/2])))
        
        
    def view(self):
        plt.figure()
        plt.imshow(self.Tdata, vmin=-1, vmax=1, extent=[0, self.traces, self.timewindow, 0])
        plt.xlabel('Trace #')
        plt.ylabel('Time [ns]')
        plt.title(self.filename)
        plt.show()

if __name__ == '__main__':
    dir='F:\\IPYdata\\RAMAC'
    date='071204'
    profileNo=4
    filename=dir+'\\'+date+'\\Profile'+str(profileNo)
    D = GPR()
    D.loadrd3(filename)
    D.subtractDC()
    D.view()